"""
1.	利用Pytorch深度学习框架，建立由卷积神经网络Resnet34构造的目标检测YOLOv1主干模型，和后端网络，
并根据给定的YOLOv1算法的loss损失函数公式，按下面要求完成相应代码。（65分）
"""
import torch as pt
import torchvision as ptv
import numpy as np


def sep(label=''):
    print('-' * 16, label, '-' * 16, sep='')


def iou(xxyy, _xxyy):
    max_x1 = max(xxyy[0], _xxyy[0])
    max_y1 = max(xxyy[1], _xxyy[1])
    min_x2 = min(xxyy[2], _xxyy[2])
    min_y2 = min(xxyy[3], _xxyy[3])
    xx = max(0., min_x2 - max_x1 + 1)
    yy = max(0., min_y2 - max_y1 + 1)
    I = xx * yy
    U = (xxyy[2] - xxyy[0] + 1) * (xxyy[3] - xxyy[1] + 1)\
        + (_xxyy[2] - _xxyy[0] + 1) * (_xxyy[3] - _xxyy[1] + 1) - I
    iou = I / (U + 1e-20)
    return iou


# ①	建立YOLOv1模型类结构
class YoloV1Model(pt.nn.Module):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # ②	导入Resnet34主干网络
        resnet34 = ptv.models.resnet34(pretrained=True)
        frontend_output_ch = resnet34.fc.in_features
        self.frontend = pt.nn.Sequential(*(list(resnet34.children())[:-2]))

        # ③	搭建后端网络模型
        self.backend = pt.nn.Sequential(
            pt.nn.Conv2d(frontend_output_ch, 1024, (3, 3), (1, 1), padding=1),
            pt.nn.BatchNorm2d(1024),
            pt.nn.ReLU(),
            pt.nn.Conv2d(1024, 1024, (3, 3), (2, 2), padding=1),
            pt.nn.BatchNorm2d(1024),
            pt.nn.ReLU(),
            pt.nn.Conv2d(1024, 1024, (3, 3), (1, 1), padding=1),
            pt.nn.BatchNorm2d(1024),
            pt.nn.ReLU(),
            pt.nn.Conv2d(1024, 1024, (3, 3), (1, 1), padding=1),
            pt.nn.BatchNorm2d(1024),
            pt.nn.ReLU(),
            pt.nn.Conv2d(1024, 30, (3, 3), (1, 1), padding=1),
            pt.nn.Sigmoid(),
        )

    def forward(self, inputs):
        # ④	打印输出YOLOv1整体模型维度结构
        print(self)
        x = self.frontend(inputs)
        print('After front end:', x.size())
        x = self.backend(x)
        print('Final size:', x.size())
        self.preds = x
        return x

    # ⑤	参考YOLOv1算法的loss损失函数计算公式
    def loss(self, labels):
        loss_proba, loss_location, loss_cls = 0., 0., 0.
        LAMBDA_COORD, LAMBDA_NONOBJ, N_GRIDS = 5, 0.5, 7
        n_smaples = len(self.preds)
        for id_sample in range(n_smaples):
            for m in range(N_GRIDS):
                for n in range(N_GRIDS):
                    # ⑥	编写目标和非目标置信度代价
                    if np.isclose(0., labels[id_sample, 4, m, n]):
                        loss_proba += LAMBDA_NONOBJ * pt.sum(self.preds[id_sample, [4, 9], m, n] ** 2)
                    else:
                        xxyy_box1 = (
                            (self.preds[id_sample, 0, m, n] + n) / N_GRIDS - self.preds[id_sample, 2, m, n] / 2,
                            (self.preds[id_sample, 1, m, n] + m) / N_GRIDS - self.preds[id_sample, 3, m, n] / 2,
                            (self.preds[id_sample, 0, m, n] + n) / N_GRIDS + self.preds[id_sample, 2, m, n] / 2,
                            (self.preds[id_sample, 1, m, n] + m) / N_GRIDS + self.preds[id_sample, 3, m, n] / 2,
                        )
                        xxyy_box2 = (
                            (self.preds[id_sample, 5, m, n] + n) / N_GRIDS - self.preds[id_sample, 7, m, n] / 2,
                            (self.preds[id_sample, 6, m, n] + m) / N_GRIDS - self.preds[id_sample, 8, m, n] / 2,
                            (self.preds[id_sample, 5, m, n] + n) / N_GRIDS + self.preds[id_sample, 7, m, n] / 2,
                            (self.preds[id_sample, 6, m, n] + m) / N_GRIDS + self.preds[id_sample, 8, m, n] / 2,
                        )
                        xxyy_label = (
                            (labels[id_sample, 0, m, n] + n) / N_GRIDS - labels[id_sample, 2, m, n] / 2,
                            (labels[id_sample, 1, m, n] + m) / N_GRIDS - labels[id_sample, 3, m, n] / 2,
                            (labels[id_sample, 0, m, n] + n) / N_GRIDS + labels[id_sample, 2, m, n] / 2,
                            (labels[id_sample, 1, m, n] + m) / N_GRIDS + labels[id_sample, 3, m, n] / 2,
                        )
                        iou1 = iou(xxyy_box1, xxyy_label)
                        iou2 = iou(xxyy_box2, xxyy_label)

                        # ⑦	编写目标边框代价
                        if iou1 > iou2:
                            loss_location += LAMBDA_COORD * pt.sum(
                                (self.preds[id_sample, 0:2, m, n] - labels[id_sample, 0:2, m, n]) ** 2\
                                + (self.preds[id_sample, 2:4, m, n].sqrt() - labels[id_sample, 2:4, m, n].sqrt()) ** 2
                            )
                            loss_proba += pt.sum((self.preds[id_sample, 4, m, n] - iou1) ** 2)
                            loss_proba += LAMBDA_NONOBJ * pt.sum((self.preds[id_sample, 9, m, n] - iou2) ** 2)
                        else:
                            loss_location += LAMBDA_COORD * pt.sum(
                                (self.preds[id_sample, 5:7, m, n] - labels[id_sample, 5:7, m, n]) ** 2\
                                + (self.preds[id_sample, 7:9, m, n].sqrt() - labels[id_sample, 7:9, m, n].sqrt()) ** 2
                            )
                            loss_proba += pt.sum((self.preds[id_sample, 9, m, n] - iou1) ** 2)
                            loss_proba += LAMBDA_NONOBJ * pt.sum((self.preds[id_sample, 4, m, n] - iou2) ** 2)

                        # ⑧	编写目标分类代价
                        loss_cls += pt.sum((self.preds[id_sample, 10:, m, n] - labels[id_sample, 10:, m, n]) ** 2)
        loss = loss_location + loss_proba + loss_cls
        loss /= n_smaples
        return loss


if '__main__' == __name__:
    model = YoloV1Model()
    x = pt.Tensor(np.zeros([4, 3, 448, 448]))
    y = pt.Tensor(np.zeros([4, 30, 7, 7]))
    y[0, 4, 0, 0] = 1.
    y[0, 4, 3, 4] = 1.
    y[2, 4, 0, 0] = 1.
    h = model(x)
    loss = model.loss(y)
    print(f'loss: {loss.detach().numpy()}')
